package cn.virens.common.mybatis.datasource;

import cn.virens.common.util.exception.APIException;
import org.springframework.jdbc.datasource.AbstractDataSource;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;

public class DynamicDataSourceImpl extends AbstractDataSource implements DynamicDataSource {
    protected final AtomicInteger counter = new AtomicInteger(0);

    protected final List<DataSource> slaveDataSource;

    protected DataSource masterDataSource;

    public DynamicDataSourceImpl(List<DataSource> slaveDataSource) {
        this.slaveDataSource = slaveDataSource;
    }

    public DynamicDataSourceImpl() {
        this(new ArrayList<>());
    }

    @Override
    public Connection getConnection(String username, String password) throws SQLException {
        return determineTargetDataSource().getConnection(username, password);
    }

    @Override
    public Connection getConnection() throws SQLException {
        return determineTargetDataSource().getConnection();
    }

    /**
     * 设置主数据库实例
     *
     * @param dataSource 数据库实例
     */
    public void setMasterDataSource(DataSource dataSource) {
        this.masterDataSource = dataSource;
    }

    /**
     * 添加从数据库实例
     *
     * @param dataSource 数据库实例
     */
    public void addSlaveDataSource(DataSource dataSource) {
        this.slaveDataSource.add(dataSource);
    }

    /**
     * 获取当前有效的数据库连接池实例
     *
     * @return 数据库连接池实例
     * @throws APIException 无
     */
    protected DataSource determineTargetDataSource() throws APIException {
        if (slaveDataSource != null && slaveDataSource.size() > 0) {
            DataSourceType type = DataSourceUtils.get0();
            if (DataSourceType.SLAVE.equals(type)) {
                return getSlaveDataSource();
            }
        }

        return masterDataSource;
    }

    /**
     * 轮询获取从数据库连接池实例
     *
     * @return 从数据库连接池实例
     * @throws APIException 无
     */
    protected DataSource getSlaveDataSource() throws APIException {
        int couid = Math.abs(counter.incrementAndGet());
        int size = Math.abs(slaveDataSource.size());

        return slaveDataSource.get(couid % size);
    }

    @Override
    public List<DataSource> getSlaveDataSources() {
        return slaveDataSource;
    }

    @Override
    public DataSource getMasterDataSource() {
        return masterDataSource;
    }
}
